package com.ai.handler;

import cn.hutool.json.JSONUtil;
import com.common.handler.RedisKeysHandler;
import jakarta.annotation.Resource;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.Set;

@Slf4j
public class RedisChatMemoryHandler implements ChatMemory {

    @Resource
    StringRedisTemplate stringRedisTemplate;

    @Data
    static class MessageBean{
        private MessageType role;
        private String content;

        public static MessageBean of(Message message){
            MessageBean messageBean = new MessageBean();
            messageBean.setRole(message.getMessageType());
            messageBean.setContent(message.getText());
            return messageBean;
        }
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        messages.forEach(message -> log.debug("Received message: " + message.getMessageType() + ", content: " + message.getText()));
        List<String> list = messages.stream()
                .filter(message -> message.getMessageType() == MessageType.USER || message.getMessageType() == MessageType.ASSISTANT)
                .map(message -> JSONUtil.toJsonStr(MessageBean.of(message))).toList();
        stringRedisTemplate.opsForList().rightPushAll(RedisKeysHandler.getAiHistoryKey(conversationId),list);
    }

    @Override
    public List<Message> get(String conversationId) {
        List<String> stringList = stringRedisTemplate.opsForList().range(RedisKeysHandler.getAiHistoryKey(conversationId), 0, -1);
        if (stringList == null || stringList.isEmpty()){
            return Collections.emptyList();
        }
        return stringList.stream().map(msg -> {
            Message message = null;
            MessageBean messageBean = JSONUtil.toBean(msg, MessageBean.class);
            switch (messageBean.getRole()){
                case USER -> message =  new UserMessage(messageBean.getContent());
                case ASSISTANT -> message =  new AssistantMessage(messageBean.getContent());
            }
            return message;
        }).filter(Objects::nonNull).toList();
    }


    @Override
    public void clear(String conversationId) {
        stringRedisTemplate.delete(RedisKeysHandler.getAiHistoryKey(conversationId));
    }
}
